from .image_convert_base import ConvertBase
import cv2 as cv
import numpy as np

class HSVConvert(ConvertBase):
    
    def __init__(self, use_rate = 0.5, is_rand = True, h = 0.5, s = 0.5, v = 0.5):
        '''
        use_rate : 使用率
        is_rand : 是否随机
        h : 色度 的增强倍率
        s : 饱和度 的增强倍率
        v : 明度 的增强倍率
        '''
        super().__init__(use_rate)
        self.is_rand = is_rand
        self._h = h
        self._s = s
        self._v = v


    
    def convert(self, img, boxes, points):
        r = np.random.uniform(-1, 1, 3) * [self._h, self._s, self._v] + 1  # random gains
        hue, sat, val = cv.split(cv.cvtColor(img, cv.COLOR_BGR2HSV))
        dtype = img.dtype  # uint8

        x = np.arange(0, 256, dtype=np.int16)
        lut_hue = ((x * r[0]) % 180).astype(dtype)
        lut_sat = np.clip(x * r[1], 0, 255).astype(dtype)
        lut_val = np.clip(x * r[2], 0, 255).astype(dtype)

        img_hsv = cv.merge((cv.LUT(hue, lut_hue), cv.LUT(sat, lut_sat), cv.LUT(val, lut_val))).astype(dtype)
        aug_img = cv.cvtColor(img_hsv, cv.COLOR_HSV2BGR)
        return aug_img, boxes, points